import cv2
import mediapipe as mp
import numpy as np

def beautify_image(image_path, output_path='output_beauty.jpg',
                   smooth_strength=0.6, whiten_strength=0.5):
    """
    :param image_path: 输入图像路径
    :param output_path: 输出图像路径
    :param smooth_strength: 磨皮强度 (0 ~ 1)
    :param whiten_strength: 美白强度 (0 ~ 1)
    """
    img = cv2.imread(image_path)
    if img is None:
        print("无法读取图片")
        return

    mp_face = mp.solutions.face_mesh
    face_mesh = mp_face.FaceMesh(static_image_mode=True, max_num_faces=1)

    img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    results = face_mesh.process(img_rgb)

    if results.multi_face_landmarks:
        # --- Step1: 磨皮（双边滤波） ---
        if smooth_strength > 0:
            d = int(10 + 20 * smooth_strength)  # 滤波直径
            sigmaColor = 50 + 100 * smooth_strength
            sigmaSpace = 50 + 100 * smooth_strength
            smoothed = cv2.bilateralFilter(img, d=d, sigmaColor=sigmaColor, sigmaSpace=sigmaSpace)
        else:
            smoothed = img.copy()

        # --- Step2: 美白（Y通道增强） ---
        if whiten_strength > 0:
            ycrcb = cv2.cvtColor(smoothed, cv2.COLOR_BGR2YCrCb)
            y, cr, cb = cv2.split(ycrcb)
            y_eq = cv2.equalizeHist(y)
            y_blend = cv2.addWeighted(y, 1 - whiten_strength, y_eq, whiten_strength, 0)
            ycrcb = cv2.merge([y_blend, cr, cb])
            whitened = cv2.cvtColor(ycrcb, cv2.COLOR_YCrCb2BGR)
        else:
            whitened = smoothed.copy()

        # --- Step3: 融合原图（保留一些原始细节） ---
        final = cv2.addWeighted(img, 1 - 0.5 * (smooth_strength + whiten_strength),
                                whitened, 0.5 * (smooth_strength + whiten_strength), 0)

        cv2.imwrite(output_path, final)
        print(f"美颜图已保存至：{output_path}")
    else:
        print("未检测到人脸")

    face_mesh.close()


if __name__ == '__main__':
    # 磨皮弱 0.2，美白强 0.8
    beautify_image("img3.jpg", "output1.jpg", smooth_strength=0.2, whiten_strength=0.8)
    # 只做美白，不磨皮
    beautify_image("img3.jpg", "output2.jpg", smooth_strength=0.0, whiten_strength=1.0)
    # 全部关闭美颜（等同于原图）
    beautify_image("img3.jpg", "output3.jpg", smooth_strength=0.0, whiten_strength=0.0)
